
import struct
import binascii
from flask_socketio import SocketIO
import zq_ota
import time
from zlib import crc32
import re

socket:SocketIO = None
is_raw_log = True

def GetChecksum(data: bytearray) -> int:
    ret = 0
    for i in range(len(data)):
        ret = (ret + (data[i] << ((i % 4) * 8)) & 0xffffffff)
    return ret

def pack_data(cmd:int, data) -> bytes:
    header = struct.pack('<BBHI', 0xAA, cmd, len(data)+8, 0)
    packet = bytearray(header + data)
    # 将校验和以小端格式填充回数据包
    checksum = (-GetChecksum(packet)) & 0xffffffff
    packet[4] = checksum & 0xFF           # 最低有效字节
    packet[5] = (checksum >> 8) & 0xFF
    packet[6] = (checksum >> 16) & 0xFF
    packet[7] = (checksum >> 24) & 0xFF   # 最高有效字节
    return bytes(packet)

def recv(data:bytearray, log1, log2):
    global fota_img_data, fota_img_type, fota_recv_ack, is_raw_log
    while len(data) >= 8:
        flag, cmd, pack_len, checksum = struct.unpack('<BBHI', data[:8])
        if flag == 0xAA:    # 找到起始字节 0xAA, 接收到带有包头的log数据
            if len(data) < pack_len:
                return data, "", ""
            
            if cmd == 0x16:      # 0x16 M0 log
                log = data[8:pack_len].decode('utf-8', errors='ignore').strip() + "\n"
                if "@Fota: ask for next chunk, offset=" in log:
                    fota_recv_ack = log
                elif "@Fota: close tmp file" in log:
                    fota_recv_ack = ""
                # socket.emit('log1', {'data': log})
                log1 += log
            elif cmd == 0x17:    # 0x17 cat1 log
                log = data[8:pack_len].decode('utf-8', errors='ignore').strip() + "\n"
                # socket.emit('log2', {'data': log})
                log2 += log
            # elif cmd == 0x10 or cmd == 0x11 or cmd == 0x12:     # 0x10/0x11/0x12 OTA pack
            #     recv_ota_ack(bytes(data[8:pack_len]))
            elif cmd == 0x09:                                   # 0x09 OTA pack
                print("recv 2...")
                fota_recv_ack = data[8:pack_len]
            data = data[pack_len:]     # 删除已处理的部分
            is_raw_log = False
        else:               # 0xAA未找到，收到原始log，解析为普通文本
            while b'\n' in data: # 查找换行符 \n 的位置
                line_end_index = data.find(b'\n')
                # 提取一行完整数据
                line_data = data[:line_end_index]
                # 删除已处理的部分
                data = data[line_end_index + 1:]
                decoded_line = line_data.decode('utf-8', errors='ignore').strip()
                # socket.emit('log1', {'data': decoded_line})
                log1 += decoded_line
            data.clear()
            # 发送要求切换到正式log命令
            # dev_cfg_pack = pack_cmd_cfg_dev()
            # dev_handle.write(dev_cfg_pack)
            is_raw_log = True
    return data, log1, log2

# 如果是原始log, 返回切换到先进log命令包, 否则返回None
def get_log_switch_pack():
    global is_raw_log
    if is_raw_log:
        return pack_data(24, bytes([(0x01 | 0x02),]))
    else:
        return None


# OTA模块 -----------------------------------------------------
fota_img_data = None
fota_img_type = 0
fota_img_start = 0

fota_recv_ack = None
fota_wait_peer = 0 # 三态0,1,2

ota_overtime = 0
# 切分组合image为m0和cat1两个部分
def unpack_combo_image(img: bytearray) -> tuple[bytearray, bytearray]:
    img_len = len(img)
    if img_len < 1024: # 合法长度肯定>=1024
        print(f"Fota image fail. data length = {img_len}")
        return None, None
    flag, ver, mcu_data_sz, mcu_data_crc, dtu_data_sz, dtu_data_crc = struct.unpack_from('<I4sIIII40x', img)
    if flag != 0x12345678: # 0x12345678开头
        print(f"Fota image fail. flag error!")
        return None, None
    if mcu_data_sz > 0:
        mcu = img[0x40:0x40+mcu_data_sz]
        if binascii.crc32(mcu) != mcu_data_crc:
            print(f"Fota image fail. mcu crc error!")
            mcu = None
    if dtu_data_sz > 0:
        dtu = img[0x40+mcu_data_sz:0x40+mcu_data_sz+dtu_data_sz]
        if binascii.crc32(dtu) != dtu_data_crc:
            print(f"Fota image fail. dtu crc error!")
            dtu = None
    return mcu, dtu

# 设置m0和cat1的升级包
def set_ota_img(img_name: str, img_data: bytearray):
    global fota_img_data, fota_img_type, fota_img_start, fota_recv_ack, fota_wait_peer
    fota_img_data = img_data
    fota_img_start = -1

    fota_recv_ack = "@Fota: ask for next chunk, offset=0"
    fota_wait_peer = 0
    if img_name.endswith('.img'):               # img文件, has m0 and cat1 parts, run with cat1 OTA way
        fota_img_type = 1
    elif img_name.endswith('.bin'):             # bin文件, only m0/cat1 img
        if fota_img_data.startswith(b'FOTA_'):  # only cat1 img, simu to run with cat1 OTA way
            fota_img_type = 1
            # 写入头部信息
            header_data = b''
            header_data += struct.pack('<I', 0x12345678) # flag
            header_data += struct.pack('<I', 0) # ver, abandoned
            # 打包M0 OTA大小和CRC（4字节无符号整数）
            header_data += struct.pack('<II', 0, 0) # m0 crc
            # 打包Cat1 OTA大小和CRC（4字节无符号整数）
            size = len(img_data) 
            crc = crc32(img_data)
            header_data += struct.pack('<II', size, crc)
            # align 64 bytes
            header_data += struct.pack('<IIIII', 0,0,0,0,0) 
            header_data += struct.pack('<IIIII', 0,0,0,0,0)
            img_data = header_data + img_data
        else:                                   # only m0 img, run with m0 OTA way
            fota_img_type = 2
    else:
        fota_img_data = None
        fota_img_type = -1
        print("不支持的文件格式\n")

def pack_cmd_mem_rw_pkg(type_val, op_val, len_val, addr_val, data_bytes):
    """
    打包ZqCmdMemRWPkg结构体
    C结构体定义:
    typedef struct {
        uint32_t type:3;        // 0 - ram; 1 - ram tmp file(<=10K); 2 - data flash; 3 - code flash; 
        uint32_t op:2;          // 0 - read; 1 - write; 2 - open, 3 - close
        uint32_t ack:3;         // 0 - no err; 1 - opened; 2 - end; 3 - busy; 4 - out of bounds; 5 - mem err; 7 - unknown error
        uint32_t res:8;         // 
        uint32_t len:16;        // read/write len
        uint32_t addr;          // read/write start addr - must align 4 bytes
        uint8_t data[0];        // data R/W or verify
    } ZqCmdMemRWPkg;
    """
    # 手动构建第一个32位字段，包含所有位域, 位域在little-endian系统中通常是从低位到高位排列
    ack_val = res_val = 0
    first_field = (
        (type_val & 0b111) |                 # bits 0-2: type
        ((op_val & 0b11) << 3) |             # bits 3-4: rw
        ((ack_val & 0b111) << 5) |           # bits 5-7: err
        ((res_val & 0b11111111) << 8) |      # bits 8-15: res
        ((len_val & 0xFFFF) << 16)           # bits 16-31: len
    )
    # 打包结构体
    # < 表示小端序
    # I 表示32位无符号整数
    # I 表示addr字段
    # {}s 表示可变长度的数据字段
    packed_data = struct.pack('<II', first_field, addr_val)
    if len_val > 0: # 读操作
        return packed_data + data_bytes
    else:
        return packed_data
    # format_str = '<II{}s'.format(len_val)
    # packed_data = struct.pack(format_str, first_field, addr_val, data_bytes)
    
    # return packed_data


def get_ota_pack() -> bytes:
    global fota_img_data, fota_img_type, fota_img_start, fota_recv_ack, fota_wait_peer

    pack = None
    # 静态变量step
    if not hasattr(get_ota_pack, 'step'):
        get_ota_pack.step = 0
        get_ota_pack.pack_len = 0

    # 没有升级包
    if fota_img_data is None:  
        return None
    # wait ack
    if fota_recv_ack is None:
        return None

    if isinstance(fota_recv_ack, str) and "@Fota: ask for next chunk, offset=" in fota_recv_ack:
        match = re.search(r'offset=(\d+)', fota_recv_ack)
        if match:
            fota_img_start = int(match.group(1))
            get_ota_pack.step = 0
            get_ota_pack.pack_len = 0
        else:
            fota_img_data = None      # 出错，结束
    elif isinstance(fota_recv_ack, bytearray) and get_cmd09_ack(fota_recv_ack) == 0:
        get_ota_pack.step += get_ota_pack.pack_len
    else:
        fota_recv_ack = None
        return None


    if fota_img_start >=0:
        start = fota_img_start + get_ota_pack.step
        end = fota_img_start + (50*1024)
        if end > len(fota_img_data): # 超出范围
            end = len(fota_img_data)
        # send normal data
        get_ota_pack.pack_len = end - start
        if get_ota_pack.pack_len > 512:
            get_ota_pack.pack_len = 512        
        data = pack_cmd_mem_rw_pkg(1, 3, get_ota_pack.pack_len, start, fota_img_data[start:(start+get_ota_pack.pack_len)]) # write tmp file
        pack = pack_data(0x09, data)

    fota_recv_ack = None
    return pack

# def recv_ota_ack(ack: bytes):
#     global fota_recv_ack
#     if ack[0] == 1:
#         fota_recv_ack = True
def get_cmd09_ack(ack: bytes):
    # 先解包前两个固定字段
    first_field, addr = struct.unpack('<II', ack[:8])
    
    # 从第一个字段中提取位域
    # type_val = first_field & 0b111                    # bits 0-2
    # op_val = (first_field >> 3) & 0b11                # bits 3-4
    ack_val = (first_field >> 5) & 0b111              # bits 5-7
    # res_val = (first_field >> 8) & 0b11111111         # bits 8-15
    # len_val = (first_field >> 16) & 0xFFFF            # bits 16-31
    return ack_val

start_time = time.time()
def DiagTimeSpan(str):
    global start_time
    end_time = time.time()  # 记录结束时间
    elapsed_time = end_time - start_time  # 计算耗时（秒）
    print(f"{str}函数耗时：{elapsed_time * 1000:.2f} ms")  # 打印耗时（毫秒）
    start_time = time.time()